#!/usr/bin/env python3
"""
模型加载和预测脚本
用于加载训练好的最佳模型并进行预测
"""

import os
import torch
import json
import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from solar_radiation_lstm import LSTMSolarModel, load_and_preprocess_data, create_sequences, SolarRadiationDataset
from torch.utils.data import DataLoader
from tqdm import tqdm
import matplotlib.pyplot as plt

def load_trained_model(model_path, training_info_path):
    """加载训练好的模型和配置信息"""
    # 加载训练信息
    with open(training_info_path, 'r', encoding='utf-8') as f:
        training_info = json.load(f)
    
    # 创建模型
    model_config = training_info['model_config']
    model = LSTMSolarModel(
        input_size=model_config['input_size'],
        hidden_size=model_config['hidden_size'],
        num_layers=model_config['num_layers'],
        output_size=model_config['output_size'],
        dropout=model_config['dropout']
    )
    
    # 加载模型参数
    model.load_state_dict(torch.load(model_path, map_location='cpu'))
    model.eval()
    
    print(f"✓ 模型加载成功: {model_path}")
    print(f"✓ 最佳验证损失: {training_info['best_val_loss']:.6f}")
    print(f"✓ 测试集RMSE: {training_info['test_rmse']:.4f}")
    
    return model, training_info

def predict_with_model(model, data_file, training_info, start_idx=0, num_predictions=10):
    """使用训练好的模型进行预测"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    
    # 加载和预处理数据
    start_year = 2022
    end_year = 2024
    data_scaled, scaler, df = load_and_preprocess_data(data_file, start_year, end_year)
    
    # 获取配置
    input_window = training_info['training_config']['input_window']
    output_window = training_info['training_config']['output_window']
    
    # 创建序列数据
    X, y = create_sequences(data_scaled, input_window, output_window)
    
    print(f"✓ 数据加载完成，总序列数: {len(X)}")
    print(f"✓ 从索引 {start_idx} 开始预测 {num_predictions} 个序列")
    
    predictions = []
    actuals = []
    
    with torch.no_grad():
        for i in tqdm(range(start_idx, min(start_idx + num_predictions, len(X))), desc="预测进度"):
            # 准备输入数据
            x = torch.FloatTensor(X[i:i+1]).to(device)  # 添加batch维度
            
            # 预测
            output = model(x)
            pred = output.cpu().numpy().flatten()
            
            # 反标准化
            pred_full = np.zeros((len(pred), scaler.n_features_in_))
            pred_full[:, 0] = pred
            pred_original = scaler.inverse_transform(pred_full)[:, 0]
            
            # 获取实际值
            actual = y[i]
            actual_full = np.zeros((len(actual), scaler.n_features_in_))
            actual_full[:, 0] = actual
            actual_original = scaler.inverse_transform(actual_full)[:, 0]
            
            predictions.append(pred_original)
            actuals.append(actual_original)
    
    return predictions, actuals, df

def plot_predictions(predictions, actuals, df, training_info, save_path=None):
    """绘制预测结果"""
    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
    
    # 计算整体评估指标
    all_preds = np.concatenate(predictions)
    all_actuals = np.concatenate(actuals)
    
    mse = np.mean((all_preds - all_actuals) ** 2)
    mae = np.mean(np.abs(all_preds - all_actuals))
    rmse = np.sqrt(mse)
    
    # 预测vs实际值散点图
    axes[0, 0].scatter(all_actuals, all_preds, alpha=0.6)
    axes[0, 0].plot([all_actuals.min(), all_actuals.max()], 
                    [all_actuals.min(), all_actuals.max()], 'r--', lw=2)
    axes[0, 0].set_xlabel('实际值 (Wh/m²)')
    axes[0, 0].set_ylabel('预测值 (Wh/m²)')
    axes[0, 0].set_title(f'预测值 vs 实际值\nRMSE: {rmse:.2f}, MAE: {mae:.2f}')
    axes[0, 0].grid(True)
    
    # 时间序列预测结果（显示前3个预测）
    for i in range(min(10, len(predictions))):
        time_steps = range(len(predictions[i]))
        axes[0, 1].plot(time_steps, actuals[i], label=f'实际值 #{i+1}', alpha=0.7)
        axes[0, 1].plot(time_steps, predictions[i], '--', label=f'预测值 #{i+1}', alpha=0.7)
    
    axes[0, 1].set_xlabel('时间步')
    axes[0, 1].set_ylabel('太阳辐照 (Wh/m²)')
    axes[0, 1].set_title('预测结果对比')
    axes[0, 1].legend()
    axes[0, 1].grid(True)
    
    # 预测误差分布
    errors = all_preds - all_actuals
    axes[1, 0].hist(errors, bins=30, alpha=0.7, edgecolor='black')
    axes[1, 0].set_xlabel('预测误差 (Wh/m²)')
    axes[1, 0].set_ylabel('频次')
    axes[1, 0].set_title(f'预测误差分布\n均值: {np.mean(errors):.2f}, 标准差: {np.std(errors):.2f}')
    axes[1, 0].grid(True)
    
    # 模型配置信息
    model_config = training_info['model_config']
    training_config = training_info['training_config']
    
    info_text = f"""模型配置:
输入窗口: {training_config['input_window']} 小时
输出窗口: {training_config['output_window']} 小时
隐藏层大小: {model_config['hidden_size']}
LSTM层数: {model_config['num_layers']}
Dropout: {model_config['dropout']}

训练结果:
最佳验证损失: {training_info['best_val_loss']:.6f}
测试集RMSE: {training_info['test_rmse']:.4f}
测试集MAE: {training_info['test_mae']:.4f}"""
    
    axes[1, 1].text(0.1, 0.9, info_text, transform=axes[1, 1].transAxes, 
                    fontsize=10, verticalalignment='top',
                    bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
    axes[1, 1].set_title('模型信息')
    axes[1, 1].axis('off')
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"预测结果图已保存到: {save_path}")
    
    plt.show()

def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='加载训练好的模型并进行预测')
    parser.add_argument('--model_path', type=str, default='solar_radiation_lstm_best_model.pth',
                        help='模型文件路径')
    parser.add_argument('--info_path', type=str, default='training_info.json',
                        help='训练信息文件路径')
    parser.add_argument('--data_file', type=str, default='POWER_Point_Hourly_20150101_20250101.csv',
                        help='数据文件路径')
    parser.add_argument('--start_idx', type=int, default=0,
                        help='开始预测的序列索引')
    parser.add_argument('--num_predictions', type=int, default=10,
                        help='预测的序列数量')
    
    args = parser.parse_args()
    
    # 检查文件是否存在
    if not os.path.exists(args.model_path):
        print(f"错误: 模型文件不存在: {args.model_path}")
        return
    
    if not os.path.exists(args.info_path):
        print(f"错误: 训练信息文件不存在: {args.info_path}")
        return
    
    if not os.path.exists(args.data_file):
        print(f"错误: 数据文件不存在: {args.data_file}")
        return
    
    # 加载模型
    model, training_info = load_trained_model(args.model_path, args.info_path)
    
    # 进行预测
    predictions, actuals, df = predict_with_model(
        model, args.data_file, training_info, args.start_idx, args.num_predictions
    )
    
    # 绘制结果
    plot_predictions(predictions, actuals, df, training_info, 'prediction_results.png')
    
    print("预测完成！")

if __name__ == "__main__":
    main() 